import pandas as pd
import numpy as np
import tushare as ts
from datetime import datetime, timedelta

# 设置tushare token（建议使用环境变量代替直接写入代码）
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()


# 获取上一个交易日
def get_last_trade_date():
    today = datetime.now()
    # 获取最近一个交易日
    trade_date = pro.trade_cal(exchange='',
                               start_date=(today - timedelta(days=7)).strftime('%Y%m%d'),
                               end_date=today.strftime('%Y%m%d'),
                               is_open='1')
    last_trade_date = trade_date.iloc[-1]['cal_date']
    return last_trade_date


# 计算RSI指标
def calculate_rsi(data, window=14):
    delta = data['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


# 获取股票历史数据并计算技术指标
def get_stock_history(ts_code, start_date, end_date):
    try:
        # 获取日线数据
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        df = df.sort_values('trade_date')

        # 计算技术指标
        df['ma5'] = df['close'].rolling(5).mean()
        df['ma10'] = df['close'].rolling(10).mean()
        df['ma20'] = df['close'].rolling(20).mean()
        df['rsi'] = calculate_rsi(df)

        # 标记买点（5日线上穿10日线）
        df['signal'] = np.where((df['ma5'] > df['ma10']) & (df['ma5'].shift(1) <= df['ma10'].shift(1)), 1, 0)

        return df
    except Exception as e:
        print(f"Error processing {ts_code}: {str(e)}")
        return None


# 主函数
def main():
    # 获取上一个交易日
    last_trade_date = get_last_trade_date()
    print(f"上一个交易日: {last_trade_date}")

    # 计算开始日期（获取足够的历史数据计算指标）
    end_date = last_trade_date
    start_date = (datetime.strptime(end_date, '%Y%m%d') - timedelta(days=60)).strftime('%Y%m%d')

    # 获取所有股票的基本信息
    stock_list = pro.stock_basic(exchange='', list_status='L')
    stock_codes = stock_list['ts_code'].tolist()

    # 存储所有买点股票
    buy_signal_stocks = []

    # 遍历股票代码（实际使用时建议分批处理或使用多线程）
    for code in stock_codes[:100]:  # 示例只处理前100只股票，实际使用时可以去掉限制
        print(f"Processing {code}...")
        hist_data = get_stock_history(code, start_date, end_date)

        if hist_data is not None and not hist_data.empty:
            # 检查最后一个交易日是否有买点信号
            last_row = hist_data.iloc[-1]
            if last_row['signal'] == 1:
                # 合并股票基本信息
                stock_info = stock_list[stock_list['ts_code'] == code].iloc[0]
                buy_signal_stocks.append({
                    'ts_code': code,
                    'name': stock_info['name'],
                    'industry': stock_info['industry'],
                    'trade_date': last_row['trade_date'],
                    'close': last_row['close'],
                    'ma5': last_row['ma5'],
                    'ma10': last_row['ma10'],
                    'rsi': last_row['rsi'],
                    'pct_chg': last_row['pct_chg']
                })

    # 转换为DataFrame
    if buy_signal_stocks:
        buy_df = pd.DataFrame(buy_signal_stocks)
        print("\n发现买点信号的股票:")
        print(buy_df[['ts_code', 'name', 'trade_date', 'close', 'ma5', 'ma10', 'rsi']])

        # 保存结果到CSV文件
        output_file = f"buy_signal_stocks_{last_trade_date}.csv"
        buy_df.to_csv(output_file, index=False, encoding='utf_8_sig')
        print(f"\n结果已保存到: {output_file}")
    else:
        print("\n今日未发现5日线上穿10日线的股票")


if __name__ == "__main__":
    main()